#include "wav_codec.h"

/* 自定义压缩文件头结构 */
typedef struct {
    unsigned int original_size;    /* 原始WAV数据大小 */
    unsigned int sample_rate;      /* 原始采样率 */
    unsigned short channels;       /* 通道数 */
    unsigned short original_bits;  /* 原始位深度 */
    unsigned int compressed_size;  /* 压缩后数据大小 */
    unsigned int target_sample_rate; /* 目标采样率 */
    unsigned short target_bits;    /* 目标位深度 */
    unsigned char use_adpcm;       /* 是否使用ADPCM编码 */
} CompressedHeader;

/* 自适应差分脉冲编码调制 (ADPCM) 的表格 */
static const int step_table[89] = {
    7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 19, 21, 23, 25, 28, 31, 34, 37, 41, 45,
    50, 55, 60, 66, 73, 80, 88, 97, 107, 118, 130, 143, 157, 173, 190, 209, 230,
    253, 279, 307, 337, 371, 408, 449, 494, 544, 598, 658, 724, 796, 876, 963,
    1060, 1166, 1282, 1411, 1552, 1707, 1878, 2066, 2272, 2499, 2749, 3024, 3327,
    3660, 4026, 4428, 4871, 5358, 5894, 6484, 7132, 7845, 8630, 9493, 10442, 11487,
    12635, 13899, 15289, 16818, 18500, 20350, 22385, 24623, 27086, 29794, 32767
};

static const int index_table[16] = {
    -1, -1, -1, -1, 2, 4, 6, 8, -1, -1, -1, -1, 2, 4, 6, 8
};

/* 默认压缩配置：保持采样率不变，降低位深度至8位，使用ADPCM */
void codec_init_config(CodecConfig* config) {
    if (config == NULL) {
        return;
    }
    
    config->target_sample_rate = 0;  /* 0将被解释为保持原采样率不变 */
    config->target_bits = 8;         /* 降低至8位 */
    config->use_adpcm = 1;           /* 使用ADPCM编码 */
}

/* 读取WAV文件头 */
static int read_wav_header(const unsigned char* buffer, unsigned int size, WAVHeader* header) {
    if (buffer == NULL || header == NULL || size < sizeof(WAVHeader)) {
        return WAV_INVALID_ARG;
    }

    memcpy(header, buffer, sizeof(WAVHeader));

    /* 验证WAV文件格式 */
    if (memcmp(header->riff_id, "RIFF", 4) != 0 || 
        memcmp(header->wave_id, "WAVE", 4) != 0 || 
        memcmp(header->fmt_id, "fmt ", 4) != 0) {
        return WAV_ERROR;
    }

    /* 仅支持PCM格式 */
    if (header->format != 1) {
        return WAV_ERROR;
    }

    return WAV_OK;
}

/* 降低采样率 */
static int downsample(const short* input, short* output, 
                     unsigned int input_samples, unsigned int factor) {
    unsigned int i, j;
    unsigned int output_samples = input_samples / factor;
    
    for (i = 0, j = 0; i < output_samples; i++, j += factor) {
        output[i] = input[j];
    }
    
    return output_samples;
}

/* 降低位深度 */
static void reduce_bit_depth(const short* input, unsigned char* output, 
                           unsigned int samples, unsigned short target_bits) {
    unsigned int i;
    int shift = 16 - target_bits;
    int max_value = (1 << target_bits) - 1;
    
    for (i = 0; i < samples; i++) {
        /* 将16位数据缩放到目标位深度 */
        int scaled = (input[i] >> shift) + (max_value / 2);
        
        /* 确保在范围内 */
        if (scaled > max_value) scaled = max_value;
        if (scaled < 0) scaled = 0;
        
        output[i] = (unsigned char)scaled;
    }
}

/* ADPCM编码一个样本 */
static unsigned char encode_adpcm_sample(short sample, short* predicted, int* index) {
    int step = step_table[*index];
    int diff = sample - *predicted;
    unsigned char code = 0;
    int tempstep;
    
    if (diff < 0) {
        code = 8;
        diff = -diff;
    }
    
    tempstep = step;
    if (diff >= tempstep) {
        code |= 4;
        diff -= tempstep;
    }
    
    tempstep >>= 1;
    if (diff >= tempstep) {
        code |= 2;
        diff -= tempstep;
    }
    
    tempstep >>= 1;
    if (diff >= tempstep) {
        code |= 1;
    }
    
    /* 调整预测值 */
    diff = step >> 3;
    if (code & 4) diff += step;
    if (code & 2) diff += step >> 1;
    if (code & 1) diff += step >> 2;
    
    if (code & 8) {
        *predicted -= diff;
        if (*predicted < -32768) *predicted = -32768;
    } else {
        *predicted += diff;
        if (*predicted > 32767) *predicted = 32767;
    }
    
    /* 调整索引 */
    *index += index_table[code & 15];
    if (*index < 0) *index = 0;
    if (*index > 88) *index = 88;
    
    return code;
}

/* ADPCM解码一个样本 */
static short decode_adpcm_sample(unsigned char code, short* predicted, int* index) {
    int step = step_table[*index];
    int diff = step >> 3;
    
    if (code & 4) diff += step;
    if (code & 2) diff += step >> 1;
    if (code & 1) diff += step >> 2;
    
    if (code & 8) {
        *predicted -= diff;
        if (*predicted < -32768) *predicted = -32768;
    } else {
        *predicted += diff;
        if (*predicted > 32767) *predicted = 32767;
    }
    
    *index += index_table[code & 15];
    if (*index < 0) *index = 0;
    if (*index > 88) *index = 88;
    
    return *predicted;
}

/* ADPCM编码音频数据 */
static int adpcm_encode(const short* input, unsigned char* output, 
                      unsigned int samples, unsigned int* out_size) {
    unsigned int i;
    unsigned int out_index = 0;
    short predicted = 0;
    int index = 0;
    
    /* 每个字节可以存储两个ADPCM样本(每个样本4位) */
    for (i = 0; i < samples; i += 2) {
        unsigned char code1 = encode_adpcm_sample(input[i], &predicted, &index);
        
        if (i + 1 < samples) {
            unsigned char code2 = encode_adpcm_sample(input[i+1], &predicted, &index);
            output[out_index++] = (code1 << 4) | code2;
        } else {
            output[out_index++] = (code1 << 4);
        }
    }
    
    *out_size = out_index;
    return WAV_OK;
}

/* ADPCM解码音频数据 */
static int adpcm_decode(const unsigned char* input, short* output, 
                      unsigned int encoded_size, unsigned int* samples) {
    unsigned int i;
    unsigned int out_index = 0;
    short predicted = 0;
    int index = 0;
    
    for (i = 0; i < encoded_size; i++) {
        unsigned char code1 = (input[i] >> 4) & 0x0F;
        output[out_index++] = decode_adpcm_sample(code1, &predicted, &index);
        
        if (out_index < *samples) {
            unsigned char code2 = input[i] & 0x0F;
            output[out_index++] = decode_adpcm_sample(code2, &predicted, &index);
        }
    }
    
    *samples = out_index;
    return WAV_OK;
}

/* PCM数据压缩（不包含文件头处理） */
int pcm_compress(const short* in_data, unsigned int in_samples,
                unsigned char* out_buffer, unsigned int* out_size,
                const CodecConfig* config) {
    unsigned int downsample_factor = 1; /* 默认不降采样 */
    short* downsampled = NULL;
    unsigned char* reduced_bits = NULL;
    unsigned int downsampled_samples;
    unsigned int adpcm_size;
    
    /* 参数检查 */
    if (in_data == NULL || out_buffer == NULL || out_size == NULL || config == NULL) {
        return WAV_INVALID_ARG;
    }
    
    /* 确定是否需要降低采样率 */
    if (config->target_sample_rate != 0 && config->target_sample_rate < 44100) {
        /* 只有当需要降低采样率时才计算降采样因子 */
        downsample_factor = 44100 / config->target_sample_rate; /* 假设原始采样率为44.1kHz */
        if (downsample_factor < 1) downsample_factor = 1;
    } else {
        /* 保持原采样率不变 */
        downsample_factor = 1;
    }
    
    /* 分配降采样或原样本缓冲区 */
    downsampled_samples = in_samples / downsample_factor;
    downsampled = (short*)malloc(downsampled_samples * sizeof(short));
    if (downsampled == NULL) {
        return WAV_NO_MEMORY;
    }
    
    /* 降低采样率或保持原样 */
    if (downsample_factor > 1) {
        downsample(in_data, downsampled, in_samples, downsample_factor);
    } else {
        /* 如果不需要降低采样率，直接复制数据 */
        memcpy(downsampled, in_data, in_samples * sizeof(short));
    }
    
    if (config->use_adpcm) {
        /* ADPCM编码 */
        int result = adpcm_encode(downsampled, out_buffer, 
                              downsampled_samples, &adpcm_size);
        
        if (result != WAV_OK) {
            free(downsampled);
            return result;
        }
        
        *out_size = adpcm_size;
    } else {
        /* 降低位深度 */
        reduced_bits = (unsigned char*)malloc(downsampled_samples);
        if (reduced_bits == NULL) {
            free(downsampled);
            return WAV_NO_MEMORY;
        }
        
        reduce_bit_depth(downsampled, reduced_bits, downsampled_samples, config->target_bits);
        
        /* 复制到输出缓冲区 */
        if (downsampled_samples > *out_size) {
            free(downsampled);
            free(reduced_bits);
            return WAV_ERROR;
        }
        
        memcpy(out_buffer, reduced_bits, downsampled_samples);
        *out_size = downsampled_samples;
        
        free(reduced_bits);
    }
    
    free(downsampled);
    return WAV_OK;
}

/* PCM数据解压缩（不包含文件头处理） */
int pcm_decompress(const unsigned char* in_buffer, unsigned int in_size,
                  short* out_data, unsigned int* out_samples,
                  unsigned char use_adpcm, unsigned short target_bits) {
    short* decompressed = NULL;
    unsigned int decompressed_samples;
    
    /* 参数检查 */
    if (in_buffer == NULL || out_data == NULL || out_samples == NULL) {
        return WAV_INVALID_ARG;
    }
    
    /* 计算解压后的样本数 */
    if (use_adpcm) {
        decompressed_samples = in_size * 2;
    } else {
        decompressed_samples = in_size;
    }
    
    /* 检查输出缓冲区大小 */
    if (*out_samples < decompressed_samples) {
        return WAV_INVALID_ARG;
    }
    
    /* 分配临时缓冲区 */
    decompressed = (short*)malloc(decompressed_samples * sizeof(short));
    if (decompressed == NULL) {
        return WAV_NO_MEMORY;
    }
    
    /* 解压数据 */
    if (use_adpcm) {
        unsigned int output_samples = decompressed_samples;
        adpcm_decode(in_buffer, decompressed, in_size, &output_samples);
        *out_samples = output_samples;
    } else {
        unsigned int i;
        
        /* 将降低位深度的数据转换回16位 */
        for (i = 0; i < decompressed_samples; i++) {
            /* 将n位数据缩放到16位 */
            int scaled = in_buffer[i];
            scaled = (scaled - (1 << (target_bits - 1))) << (16 - target_bits);
            decompressed[i] = (short)scaled;
        }
        
        *out_samples = decompressed_samples;
    }
    
    /* 复制到输出缓冲区 */
    memcpy(out_data, decompressed, *out_samples * sizeof(short));
    
    free(decompressed);
    return WAV_OK;
}

/* 压缩WAV数据 */
int wav_compress(const unsigned char* in_buffer, unsigned int in_size,
                unsigned char* out_buffer, unsigned int* out_size,
                const CodecConfig* config) {
    WAVHeader header;
    CompressedHeader comp_header;
    int result;
    unsigned int data_offset;
    const short* audio_data;
    unsigned int samples_per_channel;
    unsigned int output_buffer_size;
    unsigned int compressed_size;
    
    /* 参数检查 */
    if (in_buffer == NULL || out_buffer == NULL || out_size == NULL || config == NULL) {
        return WAV_INVALID_ARG;
    }
    
    if (*out_size < sizeof(CompressedHeader)) {
        return WAV_INVALID_ARG;
    }
    
    /* 解析WAV头 */
    result = read_wav_header(in_buffer, in_size, &header);
    if (result != WAV_OK) {
        return result;
    }
    
    /* 初始化压缩头 */
    memset(&comp_header, 0, sizeof(CompressedHeader));
    comp_header.original_size = in_size;
    comp_header.sample_rate = header.sample_rate;
    comp_header.channels = header.channels;
    comp_header.original_bits = header.bits_per_sample;
    
    /* 确定目标采样率 */
    if (config->target_sample_rate == 0) {
        comp_header.target_sample_rate = header.sample_rate; /* 保持原采样率不变 */
    } else {
        comp_header.target_sample_rate = config->target_sample_rate;
    }
    
    comp_header.target_bits = config->target_bits;
    comp_header.use_adpcm = config->use_adpcm;
    
    /* 计算数据偏移量和样本数 */
    data_offset = sizeof(WAVHeader);
    audio_data = (const short*)(in_buffer + data_offset);
    samples_per_channel = header.data_size / (header.bits_per_sample / 8) / header.channels;
    
    /* 计算剩余缓冲区大小 */
    output_buffer_size = *out_size - sizeof(CompressedHeader);
    compressed_size = output_buffer_size;
    
    /* 使用PCM压缩函数处理音频数据 */
    result = pcm_compress(audio_data, samples_per_channel * header.channels,
                        out_buffer + sizeof(CompressedHeader), &compressed_size,
                        config);
    
    if (result != WAV_OK) {
        return result;
    }
    
    comp_header.compressed_size = compressed_size;
    
    /* 写入压缩头 */
    memcpy(out_buffer, &comp_header, sizeof(CompressedHeader));
    *out_size = sizeof(CompressedHeader) + comp_header.compressed_size;
    
    return WAV_OK;
}

/* 解压缩WAV数据 */
int wav_decompress(const unsigned char* in_buffer, unsigned int in_size,
                  unsigned char* out_buffer, unsigned int* out_size) {
    CompressedHeader comp_header;
    WAVHeader wav_header;
    unsigned int data_size;
    unsigned int samples_per_channel;
    short* decompressed = NULL;
    unsigned int decompressed_samples;
    unsigned int actual_samples;
    int result;
    
    /* 参数检查 */
    if (in_buffer == NULL || out_buffer == NULL || out_size == NULL) {
        return WAV_INVALID_ARG;
    }
    
    if (in_size < sizeof(CompressedHeader)) {
        return WAV_INVALID_ARG;
    }
    
    /* 读取压缩头 */
    memcpy(&comp_header, in_buffer, sizeof(CompressedHeader));
    
    /* 计算解压后的样本数 */
    if (comp_header.use_adpcm) {
        decompressed_samples = comp_header.compressed_size * 2;
    } else {
        decompressed_samples = comp_header.compressed_size;
    }
    
    /* 检查输出缓冲区大小 */
    if (*out_size < sizeof(WAVHeader) + decompressed_samples * sizeof(short)) {
        return WAV_INVALID_ARG;
    }
    
    /* 分配解压缓冲区 */
    decompressed = (short*)malloc(decompressed_samples * sizeof(short));
    if (decompressed == NULL) {
        return WAV_NO_MEMORY;
    }
    
    /* 使用PCM解压缩函数解压数据 */
    actual_samples = decompressed_samples;
    result = pcm_decompress(in_buffer + sizeof(CompressedHeader), comp_header.compressed_size,
                          decompressed, &actual_samples,
                          comp_header.use_adpcm, comp_header.target_bits);
    
    if (result != WAV_OK) {
        free(decompressed);
        return result;
    }
    
    /* 创建WAV头 */
    memset(&wav_header, 0, sizeof(WAVHeader));
    memcpy(wav_header.riff_id, "RIFF", 4);
    memcpy(wav_header.wave_id, "WAVE", 4);
    memcpy(wav_header.fmt_id, "fmt ", 4);
    memcpy(wav_header.data_id, "data", 4);
    
    wav_header.fmt_size = 16;
    wav_header.format = 1; /* PCM */
    wav_header.channels = comp_header.channels;
    wav_header.sample_rate = comp_header.target_sample_rate;
    wav_header.bits_per_sample = 16; /* 解压后恢复为16位 */
    
    /* 计算其他WAV头字段 */
    data_size = actual_samples * sizeof(short);
    wav_header.data_size = data_size;
    wav_header.byte_rate = wav_header.sample_rate * wav_header.channels * (wav_header.bits_per_sample / 8);
    wav_header.block_align = wav_header.channels * (wav_header.bits_per_sample / 8);
    wav_header.size = 36 + data_size;
    
    /* 写入WAV头 */
    memcpy(out_buffer, &wav_header, sizeof(WAVHeader));
    
    /* 写入音频数据 */
    memcpy(out_buffer + sizeof(WAVHeader), decompressed, data_size);
    
    /* 设置输出大小 */
    *out_size = sizeof(WAVHeader) + data_size;
    
    free(decompressed);
    return WAV_OK;
} 